# uncompyle6 version 3.2.3
# Python bytecode 3.6 (3379)
# Decompiled from: Python 3.6.8 |Anaconda custom (64-bit)| (default, Feb 21 2019, 18:30:04) [MSC v.1916 64 bit (AMD64)]
# Embedded file name: api\handlers\uploads\targets_upload_handler.py
import os.path, logging
from sqlalchemy import select as sql_select
from sqlalchemy import func as sql_func
from api.handlers.uploads.base_upload_handler import UploadHandler
from db import Connection
from db.tables.uploads import UploadRow, UploadsTable
from helpers.hashing import sha256
from os.path import splitext
from settings_provider import settings
from helpers.setup_types import *

_logger = logging.getLogger("api.handlers.upload.target")


class TargetsUploadHandler(UploadHandler):
    def __init__(self, file_storage, action_handler=None):
        super().__init__()
        self.file_storage = file_storage
        self.action_handler = action_handler

    def get_descriptor(self, upload_id):
        query = (
            sql_select((UploadRow,))
            .where(UploadRow.upload_id == upload_id)
            .where(UploadRow.deleted_at.is_(None))
            .where(UploadRow.expires > sql_func.current_timestamp())
        )
        with Connection("master") as (connection):
            upload = connection.execute(query).fetchone()
            if upload is None:
                return
            upload = dict(upload)
            upload["destination"] = os.path.join(self.file_storage, upload_id)
            return upload

    def add_chunk(self, upload_id, previous_size, current_size):
        query = (
            (UploadsTable.update(values=dict(current_size=current_size)))
            .where(UploadRow.upload_id == upload_id)
            .where(UploadRow.current_size == previous_size)
            .where(UploadRow.deleted_at.is_(None))
            .where(UploadRow.status.isnot(True))
            .where(UploadRow.expires > sql_func.current_timestamp())
        )
        with Connection("master") as (connection):
            upload = connection.execute(query)
            if upload is None:
                return

    @staticmethod
    def on_success(upload_id, descriptor):
        uploads_salt = is_setup_type_on_premise() and settings.get("uploads_salt")
        if uploads_salt:
            try:
                uploads_secret = sha256(
                    descriptor["owner_id"] + upload_id, uploads_salt
                )
                from helpers.encryption.file_aes_cbc import encrypt_file

                encrypt_file(
                    descriptor["destination"],
                    descriptor["destination"] + ".enc",
                    uploads_secret,
                )
                os.remove(descriptor["destination"])
                os.rename(descriptor["destination"] + ".enc", descriptor["destination"])
            except Exception as e:
                _logger.exception("encrypting lsr failed with %s", e)

        else:
            if is_setup_type_aws():
                from helpers.target.uploads import upload_file_task_aws

                _, ext = splitext(descriptor["name"])
                destination = "%s/%s/%s" % (
                    settings["storage.targets.s3.dir"],
                    descriptor["owner_id"],
                    upload_id + ext,
                )
                settings["rq_s3_uploads"].enqueue(
                    upload_file_task_aws,
                    bucket=settings["storage.targets.s3.bucket"],
                    destination=destination,
                    source_file=descriptor["destination"],
                    secret=sha256(
                        descriptor["owner_id"] + upload_id, settings["uploads_salt"]
                    ),
                )

    def on_failure(self, upload_id, descriptor=None):
        try:
            if descriptor is None:
                descriptor = self.get_descriptor(upload_id)
            os.remove(descriptor["destination"])
        except Exception:
            pass

    def set_upload_complete(self, upload_id, descriptor=None):
        query = (
            (UploadsTable.update(values=dict(status=True, current_size=UploadRow.size)))
            .where(UploadRow.upload_id == upload_id)
            .where(UploadRow.deleted_at.is_(None))
            .where(UploadRow.status.isnot(True))
            .where(UploadRow.expires > sql_func.current_timestamp())
            .returning(UploadRow.owner_id)
        )
        with Connection("master") as (connection):
            owner_id = connection.execute(query).scalar()
            if owner_id is None:
                self.on_failure(upload_id, descriptor)
                return
        if not self.action_handler:
            self.on_success(upload_id, descriptor)
            return
        actions = descriptor["actions"]
        if not actions:
            self.on_success(upload_id, descriptor)
            return
        actions = actions.get("ok")
        if not actions:
            self.on_success(upload_id, descriptor)
            return
        for action in actions:
            method_name, kwargs = action
            method = getattr(self.action_handler, method_name)
            if kwargs is None:
                kwargs = dict()
            kwargs["upload_descriptor"] = descriptor
            method(**kwargs)

        self.on_success(upload_id, descriptor)
